'''
    处理3列数据的绘图模块
        1.绘图函数后加_layout为相应的风格设置函数
'''
import plotly.graph_objs as go
import pandas as pd
import numpy as np
from pandas.api.types import is_string_dtype

# 第三列为数值,绘制3d-scatter
def three_dim_scatter_z_num(df,col1, col2, col3):
    z = df[col3]
    if(is_string_dtype(z)):
        z=pd.factorize(z)[0].astype(np.uint16)
    trace = go.Scatter3d(
        x=df[col1], y=df[col2], z=df[col3],
        mode='markers', marker={'size': 8, 'color': z, 'colorscale': 'Blackbody', 'opacity': 0.8, "showscale": True,
                                "colorbar": {"thickness": 15, "len": 0.5, "x": 0.8, "y": 0.6, }, })
    return trace

# 3d-scatter外观
def three_dim_scatter_layout(col1,col2,col3):
    title=','.join([col1,col2,col3])
    layout=go.Layout(
                height=700, title=title,
                paper_bgcolor="#f3f3f3",
                scene={"aspectmode": "cube", "xaxis": {"title":col1, },
                       "yaxis": {"title": col2, },
                       "zaxis": {"title": col3, }})
    return layout

# 第三列为字符串,绘制3d-scatter
def three_dim_scatter_z_str(df,col1, col2, col3):
    z = pd.factorize(df[col3])[0].astype(np.uint16)
    trace = go.Scatter3d(
        x=df[col1], y=df[col2], z=df[col3],
        mode='markers', marker={'size': 8, 'color': z, 'colorscale': 'Blackbody', 'opacity': 0.8, "showscale": False,
                                "colorbar": {"thickness": 15, "len": 0.5, "x": 0.8, "y": 0.6, }, })
    return trace